import enum
import os
from re import T

import shutil
try:
    from .cvio import cvio
    from .labelme2coco import gen_ann_key_info
except:
    from cvio import cvio
    from labelme2coco import gen_ann_key_info    

def parse_text(text):
    if os.path.isfile(text):
        assert text.endswith('.txt')
        with open(text) as fp:
            labels = [x.strip().replace('\t', '').replace('\n', '') for x in fp.readlines()]
        return labels
    labels = text.replace('\t', ',').replace('\n', ',').replace('，', ',').replace(' ', ',').split(',')
    return labels

def collect_images_with_label(src, label, merge=False, mode='copy', dst=''):
    imgann = cvio.load_img_ann_list(src,recursive=False)

    n = len(imgann)

    move = shutil.copy
    targets = parse_text(label)
    # targets = label.replace(' ', ',').replace('，',',').replace('\t', ',').replace('\n', ',').replace('   ', ',').split(',')
    for i, (img, ann) in enumerate(imgann, 1):
        print('[%d/%d] %s' % (i, n, os.path.basename(img)))
        ann_info = cvio.load_ann(ann)
        labels = [s['label'] for s in ann_info['shapes']]
        for target in targets:
            if target == '':
                continue
            if target in labels:
                if merge and dst not in (src, "", None):
                    _dst = dst
                else:
                    _dst = os.path.join(dst, target) if dst not in (None, '') else os.path.join(src, target)
                if not os.path.exists(_dst):
                    os.makedirs(_dst)
                _dst_img = os.path.join(_dst, os.path.basename(img))
                if not os.path.exists(_dst_img):
                    move(img, _dst_img)
                _dst_ann = os.path.join(_dst, os.path.basename(ann))
                if not os.path.exists(_dst_ann):
                    move(ann, _dst_ann)

def checkWrongMasks(src):
    imgann = cvio.load_img_ann_list(src,recursive=False)
    num = len(imgann)
    print(num, 'images')

    fp = open(os.path.join(src, 'wrong.txt'), 'w')
    for i, (img, ann) in enumerate(imgann, 1):
        print('[%d/%d]' % (i, num), img)
        
        nimpath = 0
        nimdata = 0
        nstype = 0
        nimwh = 0
        npoints = 0        

        annInfo = cvio.load_ann(ann)
        w = annInfo['imageWidth']
        h = annInfo['imageHeight']
        obn = os.path.basename(img)
        abn = annInfo['imagePath']

        flag = False

        if obn != abn:
            flag = True
            nimpath += 1
            annInfo['imagePath'] = obn

        if annInfo['imageData']:
            annInfo['imageData'] = None
            nimdata += 1
            flag = True

        imarray = cvio.load_img(img, mode='cv2')
        oh, ow = imarray.shape[:2]                
        if oh != h or ow != w:
            nimwh += 1
            annInfo['imageHeight'] = oh
            annInfo['imageWidth'] = ow
            flag = True

        shapes = annInfo['shapes']
        shapes_ = []
        if len(shapes):
            for shape in shapes:
                label = shape['label']
                points = shape['points']

                if len(points) == 0:
                    npoints += 1
                    flag = True
                    continue 

                _, bbox, area = gen_ann_key_info(points)
                if area <= 1:
                    npoints += 1
                    flag = True
                    continue
                if len(points) == 2 and shape['shape_type'] != 'rectangle':
                    nstype += 1
                    shape['shape_type'] = 'rectangle'
                    flag = True
                if len(points) > 2 and shape['shape_type'] == 'rectangle':
                    nstype += 1
                    shape['shape_type'] = 'polygon'
                    flag = True                    
                _points = []
                for p in points:
                    x, y = p
                    if x < 0:
                        flag = True
                        npoints += 1
                        x = 0
                    if y < 0:
                        flag = True
                        npoints += 1
                        y = 0
                    if x >= ow:
                        npoints += 1
                        flag = True
                        x = w - 1
                    if y >= oh:
                        npoints += 1
                        flag = True                        
                        y = h - 1
                    p = [x, y]
                    _points += [[x, y]]
                shape['points'] = _points
                shapes_.append(shape)

        if flag:
            print('*****' * 20)
            print('imagePath %d imageData %d imageWH %d shape_type %d  points %d ' % (nimpath, nimdata, nimwh, nstype, npoints))
            print('*****' * 20)
            annInfo['shapes'] = shapes_
            cvio.write_ann(annInfo, ann)
            print('%s ERROR_TYPE imagePath %d imageData %d imageWH %d shape_type %d  points %d ' % (obn, nimpath, nimdata, nimwh, nstype, npoints), file=fp)
    fp.close()


if __name__ == '__main__':
    checkWrongMasks(r'G:\data\datasets\drink\daiding_101\train\mask')            